#include "stdafx.h"
#include "SocketBase.h"
#include <WS2tcpip.h>

bool CSocketBase::CreateSocket(SOCKET & sock)
{
	sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (INVALID_SOCKET == sock)
		return false;
	else
		return true;
}

CSocketBase::CSocketBase(WORD port) :m_iListenPort(port), m_ListenSocket(INVALID_SOCKET)
{
	WSADATA wsa = { 0 };
	if (WSAStartup(MAKEWORD(2, 2), &wsa) == SOCKET_ERROR) {
		m_ListenSocket = INVALID_SOCKET;
		WSACleanup();
	}
}

CSocketBase::~CSocketBase()
{
}

BOOL CSocketBase::Listen()
{


	SOCKADDR_IN srvAddr = { 0 };
	srvAddr.sin_family = AF_INET;
	srvAddr.sin_addr.S_un.S_addr = htonl(INADDR_ANY);
	srvAddr.sin_port = htons(m_iListenPort);

	if ((m_ListenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) == INVALID_SOCKET) {
		return FALSE;
	}

	if (bind(m_ListenSocket, (SOCKADDR*)&srvAddr, sizeof(srvAddr)) == INVALID_SOCKET) {
		Close();
		return FALSE;
	}

	if (listen(m_ListenSocket, SOMAXCONN) == SOCKET_ERROR) {
		Close();
		return FALSE;
	}
	return TRUE;
}

void CSocketBase::Accept(SOCKET &ClientSocket)
{
	ClientSocket = INVALID_SOCKET;
	sockaddr_in clientAddress = { 0 };
	int addrLen = sizeof(clientAddress);
	ClientSocket = accept(m_ListenSocket, (SOCKADDR*)&clientAddress, &addrLen);
}

int CSocketBase::ReceiveFrom(const SOCKET sock, char * szBuffer, int iBufferSize)
{
	return recv(sock, szBuffer, iBufferSize, 0);
}

int CSocketBase::SendTo(const SOCKET sock, const char * szBuffer, int iBufferSize)
{
	int pos = 0;
	while (pos < iBufferSize)
	{
		int ret = send(sock, szBuffer + pos, iBufferSize - pos, 0);

		if (ret > 0)
		{
			pos += ret;
		}
		else
		{
			return ret;
		}
	}
	return pos;
}

BOOL CSocketBase::ConnectToServer(SOCKET & sockfd, const char * szHostName, WORD iPort)
{	
	// code for a client connecting to a server
	// namely a stream socket to www.example.com on port 80 (http)
	// either IPv4 or IPv6
	BOOL bConnected = FALSE;
	struct addrinfo hints, *servinfo, *p;

	int rv;

	memset(&hints, 0, sizeof hints);
	hints.ai_family = AF_INET; // use AF_INET6 to force IPv6
	hints.ai_socktype = SOCK_STREAM;

	if ((rv = getaddrinfo(szHostName, "", &hints, &servinfo)) != 0) {
		//fPRINT(stderr, "getaddrinfo: %s\n", gai_strerror(rv));
		//exit(1);
		bConnected = FALSE;
	}
	else {
		// loop through all the results and connect to the first we can
		for (p = servinfo; p != NULL; p = p->ai_next) {
			if (p->ai_family == AF_INET) {
				if (p->ai_socktype == SOCK_STREAM) {
					//Create socket

					if (!CreateSocket(sockfd))continue;


					//Get ip
					IN_ADDR addr = ((SOCKADDR_IN*)p->ai_addr)->sin_addr;
					char szIP[INET_ADDRSTRLEN] = { 0 };
					inet_ntop(AF_INET, &addr, szIP, INET_ADDRSTRLEN);
					//Connect
					SOCKADDR_IN srvAddr = { 0 };
					srvAddr.sin_family = AF_INET;
					srvAddr.sin_addr = addr;
					srvAddr.sin_port = htons(iPort);

					//inet_pton(AF_INET, szIP, &srvAddr.sin_addr);

					if (connect(sockfd, (SOCKADDR*)&srvAddr, sizeof(srvAddr)) == SOCKET_ERROR) {
						DWORD d = GetLastError();
						closesocket(sockfd);
						sockfd = INVALID_SOCKET;
						continue;
					}

					bConnected = TRUE;

					//PRINT("Domain:%s IP%s\n",szHostName, szIP);
					break; // if we get here, we must have connected successfully
				}
			}
		}
		if (p == NULL) {
			// looped off the end of the list with no connection
			//fPRINT(stderr, "failed to connect\n");
			//exit(2);
			bConnected = FALSE;
		}
	}



	freeaddrinfo(servinfo); // all done with this structure

	return bConnected;
}

void CSocketBase::Close()
{
	if (INVALID_SOCKET != m_ListenSocket) {
		shutdown(m_ListenSocket, SD_BOTH);
		closesocket(m_ListenSocket);
		m_ListenSocket = INVALID_SOCKET;
	}
}

void CSocketBase::Close(SOCKET & sock)
{
	if (INVALID_SOCKET != sock) {
		shutdown(sock, SD_BOTH);
		closesocket(sock);
		sock = INVALID_SOCKET;
	}
}

string CSocketBase::getIP(SOCKET & sock)
{
	SOCKADDR addr = { 0 };
	SOCKADDR_IN * addr_v4 = { 0 };
	int addrLen = sizeof(addr);

	string ip;

	if (0 == getpeername(sock, &addr, &addrLen)) {
		if (addr.sa_family == AF_INET) {
			addr_v4 = (SOCKADDR_IN*)&addr;
			char szBuffer[INET_ADDRSTRLEN];
			inet_ntop(AF_INET, (void*)&addr_v4->sin_addr, szBuffer, INET_ADDRSTRLEN);
			ip = szBuffer;
		}
	}
	return ip;
}
